#include <stdlib.h>
#include <string.h>
#include <pcap.h>
#include <netinet/ether.h>
#include <netinet/ip.h>
#include <netinet/udp.h>
#include <netinet/tcp.h>
#include <assert.h>

#include "net.h"

/* support for -iany tcpdump switch */
/* copied from tcpdump sources */
#define SLL_HDR_LEN 16
#define SLL_ADDRLEN 8

struct sll_header {
  u_int16_t sll_pkttype;
  u_int16_t sll_hatype;
  u_int16_t sll_halen;
  u_int8_t sll_addr[SLL_ADDRLEN];
  u_int16_t sll_protocol;
};

/* end of copied from tcpdump sources */

static flow_t *_flowset_lookup(flowset_t *set, uint8_t protocol, uint32_t src,
  uint16_t source, uint32_t dst, uint16_t dest) {
  flow_t *flow;
  for (flow = set->flows.tqh_first; flow; flow = flow->next.tqe_next) {
    if (flow->protocol == protocol &&
      (flow->connector.sin_addr.s_addr == dst &&
      flow->connector.sin_port == dest) &&
      (flow->connectee.sin_addr.s_addr == src &&
      flow->connectee.sin_port == source))
      return flow;
    if (flow->protocol == protocol &&
      (flow->connectee.sin_addr.s_addr == dst &&
      flow->connectee.sin_port == dest) &&
      (flow->connector.sin_addr.s_addr == src &&
      flow->connector.sin_port == source))
      return flow;
  }
  return NULL;
}

TAILQ_HEAD(flowq,flow);

static flow_t *_flowset_insert(flowset_t *set, uint8_t protocol, uint32_t src,
  uint16_t source, uint32_t dst, uint16_t dest) {
  flow_t *flow, *previous_flow;
  flow = calloc(1, sizeof(*flow));
  assert(flow);
  flow->connector.sin_addr.s_addr = src;
  flow->connector.sin_port = source;
  flow->connectee.sin_addr.s_addr = dst;
  flow->connectee.sin_port = dest;
  flow->protocol = protocol;
  flow->type = MONOLOG_FLOW;
  flow->last_packet = NULL;
  TAILQ_INIT(&flow->packets);
  /* odds are that flow should be the last of the list */
  TAILQ_FOREACH_REVERSE(previous_flow,&set->flows,flowq,next) {
    if (timercmp(&previous_flow->when,&flow->when,<))
      break;
  }
  if (!previous_flow)
    TAILQ_INSERT_TAIL(&set->flows,flow,next);
  else
    TAILQ_INSERT_AFTER(&set->flows,previous_flow,flow,next);
  set->n_flows++;
  return flow;
}

TAILQ_HEAD(packetq,packet);

static packet_t *_flow_insert(flow_t *flow, flowset_t *set, uint32_t src, uint16_t source,
  uint32_t dst, uint16_t dest, const char *str, size_t size) {
  packet_t *packet, *previous_packet;
  packet = calloc(1, sizeof(*packet));
  assert(packet);
  if ((flow->connector.sin_addr.s_addr == src &&
    flow->connector.sin_port == source) &&
    (flow->connectee.sin_addr.s_addr == dst &&
    flow->connectee.sin_port == dest))
    packet->type = REQUEST_PACKET;
  if ((flow->connectee.sin_addr.s_addr == src &&
    flow->connectee.sin_port == source) &&
    (flow->connector.sin_addr.s_addr == dst &&
    flow->connector.sin_port == dest))
    packet->type = RESPONSE_PACKET;
  packet->segment.iov_base = malloc(size);
  assert(packet->segment.iov_base);
  packet->segment.iov_len = size;
  memcpy(packet->segment.iov_base, str, size);
  packet->flow = flow;
  /* sort in time order */
  TAILQ_FOREACH_REVERSE(previous_packet,&flow->packets,packetq,next_by_flow) {
    if (timercmp(&previous_packet->when,&packet->when,<))
      break;
  }
  if (!previous_packet)
    TAILQ_INSERT_TAIL(&flow->packets,packet,next_by_flow);
  else
    TAILQ_INSERT_AFTER(&flow->packets,previous_packet,packet,next_by_flow);
  /* the same */
  TAILQ_FOREACH_REVERSE(previous_packet,&set->packets,packetq,next_by_time) {
    if (timercmp(&previous_packet->when,&packet->when,<))
      break;
  }
  if (!previous_packet)
    TAILQ_INSERT_TAIL(&set->packets,packet,next_by_time);
  else
    TAILQ_INSERT_AFTER(&set->packets,previous_packet,packet,next_by_time);
  if (flow->last_packet && *flow->last_packet != packet->type) {
    flow->type = DIALOG_FLOW;
  }
  flow->last_packet = &packet->type;
  return packet;
}

struct fragment;
struct bag {
  flowset_t *set;
  int link;
  TAILQ_HEAD(,fragment) fragments;
};

struct fragment {
  uint32_t src;
  uint32_t dst;
  uint16_t id;
  uint8_t protocol;
  struct timeval when;
  struct iovec segment;
  TAILQ_ENTRY(fragment) next;
};

static struct fragment *_fragment_insert(struct bag *bag, uint32_t src,
  uint32_t dst, uint16_t id, uint8_t protocol) {
  struct fragment *fragment;
  fragment = calloc(1, sizeof(*fragment));
  assert(fragment);
  fragment->src = src;
  fragment->dst = dst;
  fragment->id = id;
  fragment->protocol = protocol;
  TAILQ_INSERT_TAIL(&bag->fragments,fragment,next);
  return fragment;
}

static void _fragment_remove(struct bag *bag, struct fragment *fragment) {
  TAILQ_REMOVE(&bag->fragments,fragment,next);
  free(fragment->segment.iov_base);
  free(fragment);
}

static struct fragment *_fragment_lookup(struct bag *bag, uint32_t src,
  uint32_t dst, uint16_t id) {
  struct fragment *fragment;
  for (fragment = bag->fragments.tqh_first; fragment;
    fragment = fragment->next.tqe_next)
    if (fragment->src == src && fragment->dst == dst
      && fragment->id == id)
      return fragment;
  return NULL;
}

static void _fragment_append(struct fragment *fragment, const char *str,
  size_t size) {
  char *tmp;
  tmp = realloc(fragment->segment.iov_base, fragment->segment.iov_len + size);
  assert(tmp);
  memcpy(tmp + fragment->segment.iov_len, str, size);
  fragment->segment.iov_base = tmp;
  fragment->segment.iov_len += size;
}

static void _pcap_callback(u_char *user, const struct pcap_pkthdr *hdr,
  const u_char *pkt) {
  /* low level headers */
  struct ether_header *eth;
  struct sll_header *sll;
  /* ipv4 header */
  struct iphdr *ip;
  /* tcp/udp headers */
  struct tcphdr *tcp;
  struct udphdr *udp;
  struct iovec segment;
  uint8_t protocol;
  uint16_t sport, dport;
  struct fragment *fragment;
  flow_t *flow;
  packet_t *packet;

  struct bag *bag = (struct bag *) user;

  /* check if captured length is length of packet */
  /* saying no here means that capture was done using default */
  /* tcpdump behaviour which truncates packet to 68 bytes */
  assert(hdr->caplen == hdr->len);

  switch (bag->link) {
  case DLT_EN10MB:
    eth = (struct ether_header *) pkt;
    /* ignore non IPv4 packets */
    if (eth->ether_type != htons(ETHERTYPE_IP))
      return;

    /* set segment values */
    segment.iov_base = (char *)pkt + sizeof(*eth);
    segment.iov_len = hdr->len - sizeof(*eth);
    break;
  case DLT_LINUX_SLL:
    sll = (struct sll_header *) pkt;
    /* ignore non IPv4 packets */
    if (sll->sll_protocol != htons(ETHERTYPE_IP))
      return;

    /* set segment values */
    segment.iov_base = (char *)pkt + sizeof(*sll);
    segment.iov_len = hdr->len - sizeof(*sll);
    break;
  default:
    /* this should not be reached ever */
    abort();
  }

  ip = (struct iphdr *) segment.iov_base;
  segment.iov_base = (char *) segment.iov_base + ((ip->ihl)<<2);
  segment.iov_len -= ((ip->ihl)<<2);

  /* IPv4 fragmentation macros */
#define IP_MOREFRAGS(IPH)	((IPH)->frag_off & IP_MF)
#define IP_ISFRAG(IPH)		(((IPH)->frag_off & (IP_MF | IP_OFFMASK)) == 0)
#define IP_FIRSTFRAG(IPH)	(((IPH)->frag_off & IP_OFFMASK) == 0)
#define IP_FRAGOFF(IPH)		(((IPH)->frag_off & IP_OFFMASK) << 3)
  /* handle fragmentation */
  fragment = NULL;
  if (IP_ISFRAG(ip)) {
    /* if not first first fragment */
    if (!IP_FIRSTFRAG(ip)) {
      fragment = _fragment_lookup(bag, ip->saddr, ip->daddr, ip->id);
      /* if previous fragment does not exist, discard packet */
      if (!fragment)
        return;
      else {
        /* check that reassembled length equals offset */
        assert(IP_FRAGOFF(ip) == fragment->segment.iov_len);

        _fragment_append(fragment, segment.iov_base, segment.iov_len);
      }
    }
    else {
      fragment = _fragment_lookup(bag, ip->saddr, ip->daddr, ip->id);
      assert(!fragment);

      fragment = _fragment_insert(bag, ip->saddr, ip->daddr, ip->id,
        ip->protocol);
      assert(fragment);

      memcpy(&fragment->when, &hdr->ts, sizeof(hdr->ts));

      _fragment_append(fragment, segment.iov_base, segment.iov_len);
    }
  }

  /* if packet is not fragmented or has been reassembled */
  if (!IP_MOREFRAGS(ip)) {
    if (fragment) {
      segment.iov_base = fragment->segment.iov_base;
      segment.iov_len = fragment->segment.iov_len;
      protocol = fragment->protocol;
    }
    else {
      protocol = ip->protocol;
    }

    sport = dport = 0;
    switch (protocol) {
    case IPPROTO_UDP:
      udp = (struct udphdr *) segment.iov_base;
      sport = udp->source;
      dport = udp->dest;

      segment.iov_base = (char *) segment.iov_base + sizeof(*udp);
      segment.iov_len -= sizeof(*udp);
      break;
    case IPPROTO_TCP:
      tcp = (struct tcphdr *) segment.iov_base;
      sport = tcp->source;
      dport = tcp->dest;

      segment.iov_base = (char *) segment.iov_base + ((tcp->doff)<<2);
      segment.iov_len -= ((tcp->doff)<<2);
      break;
    }

    flow = _flowset_lookup(bag->set, protocol, ip->saddr, sport, ip->daddr,
      dport);
    if (!flow) {
      flow = _flowset_insert(bag->set, protocol, ip->saddr, sport, ip->daddr,
        dport);
      if (bag->link == DLT_LINUX_SLL && sll->sll_halen != 0) {
        switch (sll->sll_pkttype) {
        case 0:
          flow->role = UNICAST_TO_US;
          break;
        case 0x0400:
          flow->role = SENT_BY_US;
          break;
        default:
          abort();
        }
      }
      memcpy(&flow->when, &hdr->ts, sizeof(hdr->ts));
    }

    if (segment.iov_len == 0)
      return;

    /* check direction and append segment to flow */
    packet = _flow_insert(flow, bag->set, ip->saddr, sport, ip->daddr, dport,
      segment.iov_base, segment.iov_len);

    if (fragment) {
      memcpy(&packet->when, &fragment->when, sizeof(fragment->when));
      _fragment_remove(bag, fragment);
    }
    else
      memcpy(&packet->when, &hdr->ts, sizeof(hdr->ts));
  }
}

flowset_t *flowset_load(const char *file, const char *filter) {
  struct bag bag;
  flowset_t *set;
  pcap_t *pcap;
  char errbuf[PCAP_ERRBUF_SIZE];
  struct bpf_program fp;
  int ret;
  static char default_filter[] = "tcp or udp";
  char *real_filter;

  pcap = pcap_open_offline(file, errbuf);
  assert(pcap);

  set = calloc(1, sizeof(*set));
  assert(set);
  TAILQ_INIT(&set->flows);
  TAILQ_INIT(&set->packets);

  bag.set = set;
  bag.link = pcap_datalink(pcap);
  assert(bag.link == DLT_EN10MB || bag.link == DLT_LINUX_SLL);
  TAILQ_INIT(&bag.fragments);

  if (filter)
    real_filter = strdup(filter);
  else
    real_filter = strdup(default_filter);
  assert(real_filter);

  ret = pcap_compile(pcap, &fp, real_filter, 0, 0);
  assert(ret != -1);

  ret = pcap_setfilter(pcap, &fp);
  assert(ret != -1);

  pcap_dispatch(pcap, -1, _pcap_callback, (u_char *) &bag);

  pcap_close(pcap);
  free(real_filter);

  return set;
}
